#light
namespace MPITools
#nowarn "0051";;  // suppress warnings about the use of byref features to interoperate with native code
#nowarn "9";;  // suppress warnings about the use of byref features to interoperate with native code

open System.Runtime.InteropServices
open Microsoft.FSharp.NativeInterop
open NativeInterop.NativePtr
open System
open Microsoft.FSharp.Math
open MPI.Bindings.MPIRaw
open System.IO


module MPI =
    let pinA arr = PinnedArray.of_array(arr)
    let freeA (pA: 'a PinnedArray) = pA.Free()
    let logging = ref false
    let private TypeConvert (t) =                  
         let res =
          match (box t) with         
            | :? byte    -> MPI_Datatype.MPI_BYTE         
            | :? char    -> MPI_Datatype.MPI_CHAR
            | :? int16   -> MPI_Datatype.MPI_SHORT
            | :? uint32  -> MPI_Datatype.MPI_LONG                
            | :? int     -> MPI_Datatype.MPI_INT            
            | :? double  -> MPI_Datatype.MPI_DOUBLE   
            | _       ->    let s =  "not implemented data type :" + (((box t).GetType()).FullName)
                            failwith s
         int res
         
    let private TTypeConvert (t : Type) =         
         let res =  if t =  typeof<byte> then MPI_Datatype.MPI_BYTE         
                    elif t= typeof<char> then MPI_Datatype.MPI_WCHAR                
                    elif t = typeof<int> then MPI_Datatype.MPI_INT          
                    elif t= typeof<float> then MPI_Datatype.MPI_FLOAT        
                    elif t= typeof<double> then MPI_Datatype.MPI_DOUBLE   
                    elif t.IsClass && t.IsSerializable then MPI_Datatype.MPI_BYTE
                    else let s =  "not implemented data type"    + t.FullName
                         failwith  s                                                  
         int res         
         
         
    type 'a  MPIType =
        | MPIChunk of (nativeint * int * int)
        | MPIChunkArray  of ('a[] * int * int)
        | MPISerializedArray of ((byte[]*int*int)[])
       
    let (|Array|Primitive|Serializable|)  (typ: Type) =        
        if typ.IsPrimitive        then Primitive    
        elif typ.IsArray            then Array(typ.GetElementType(),typ.GetArrayRank())
        elif typ.IsSerializable     then Serializable   
        else failwith               "Not supported type for communication: %s" typ.Name         
        
    let private failonError result failFunction =
        if result <> MPI_SUCCESS then 
            MPI_Abort(MPI_COMM_WORLD, result) |> ignore
            failwith (failFunction ^ " function exited with code: "^ result.ToString())
        
    let initialize() =     
        let args = System.Environment.GetCommandLineArgs()
        let mutable argsNum = args.Length                               
        let mutable args  = args |> 
                            Array.map (fun arg ->
                            NativeInterop.NativePtr.of_nativeint (Marshal.StringToHGlobalAnsi(arg)) :char nativeptr)
        
        let mutable m_args = pinA args   
        let mutable argPtr = m_args.Ptr                  
        try failonError (MPI_Init(&&argsNum, &&argPtr))  "MPI_Init"            
        finally freeA m_args
                                    
    let finalize() = 
        failonError (MPI_Finalize())   "MPI_Finalize"

    let size() =
        let mutable mysize = -1
        failonError (MPI_Comm_size(MPI_COMM_WORLD, &&mysize)) "size" 
        mysize
                        
    let rank() =          
       let mutable myrank = -1
       failonError  (MPI_Comm_rank(MPI_COMM_WORLD, &&myrank))   "MPI_Comm_rank"
       myrank    

    let private serialize data =                 
            use memStream = new MemoryStream()
            let serializer = new System.Runtime.Serialization.Formatters.Binary.BinaryFormatter()
            serializer.Serialize(memStream,data)
            let buffer = memStream.GetBuffer()
            memStream.Close()
            buffer
            
    let private deserialize (buffer:byte[]) = 
        use memoryStream = new System.IO.MemoryStream(buffer)
        let deserializer = new System.Runtime.Serialization.Formatters.Binary.BinaryFormatter()
        let buffer = deserializer.Deserialize(memoryStream)
        memoryStream.Close()
        buffer        
    
    let private dataprepare arr= let mutable arr = arr
                                 let mutable pinnedbuffer = pinA arr
                                 ((pinnedbuffer.Ptr |> to_nativeint), pinnedbuffer)         
   
    let  private log sdata datatype =
        if !logging then printf "sdata Type:%s MPIType:%d\n" (((box sdata).GetType()).FullName) datatype
        
    let private logR destination tag = 
        if !logging then printf "Destination:%d Tag:%d\n" destination tag
        
    let private mpisend sdata item   datatype destination tag =   
        let mutable sdata = sdata   
        log sdata datatype       
        failonError (MPI_Send( (&&sdata |>  to_nativeint),item, datatype, destination, tag, MPI_COMM_WORLD)) "MPI_Send" 
        logR destination tag
    
    let private mpireceive<'a>    source tag = 
        let mutable status = new MPI_Status()
        let mutable rdata = ((Array.zero_create 1) :'a array).[0] //something like that in IL -> initobj    !T         
        log rdata 1 
        failonError (MPI_Recv((&&rdata |>  to_nativeint), 1, (TypeConvert  rdata), source, tag,MPI_COMM_WORLD, &&status)) "MPI_Recv"
        logR source tag
        rdata    

    // it sends the size first than the actual array
    let private mpiarrsend (sdata:'a array) items datatype destination tag =        
            mpisend items 1 (TypeConvert items) destination tag

            let datatosend, freearr = dataprepare sdata
            let res =try log sdata datatype
                         logR destination tag
                         MPI_Send(datatosend,items, datatype, destination, tag, MPI_COMM_WORLD)
                     finally freeA  freearr
            
            failonError res "MPI_Send"
            
        
    let private mpiarrayreceive<'a>  source tag = 
        let mutable status = new MPI_Status()
        let mutable size = mpireceive<int>  source tag 
        if size <= 0 then failwith "size can't be 0"
        let mutable buffer = Array.zero_create size : 'a array
        
        let datatoreceive, freearr = dataprepare buffer  
        let res = 
            try log buffer 1
                logR source tag  
                MPI_Recv(datatoreceive, size, (TTypeConvert (typeof<'a>) ), source, tag,MPI_COMM_WORLD, &&status)
            finally  freeA freearr        
        failonError res "MPI_Recv"        
        buffer            
                                               
     
    let send(data,destination, tag) =             
        match (box data).GetType() with
        | Primitive       ->  mpisend data 1  (TypeConvert data) destination tag                                      
        | Serializable _ -> let buffer= serialize data                            
                            mpiarrsend (buffer) (buffer.Length) (int MPI_Datatype.MPI_BYTE) destination tag                                                                                   
        | _ -> failwith "Not supported from send function"    
        
    let receive<'a>(source,tag)  =         
        match typeof<'a> with
        | Primitive ->  mpireceive<'a>   source tag                          
   
        | Serializable -> let buffer = mpiarrayreceive<byte>  source tag   
                          use memStream = new MemoryStream(buffer)
                          let desObj = deserialize buffer
                          memStream.Close(); 
                          desObj :?> 'a        
        | _ -> failwith "Not supported from receive function"            
        
    let sendArray(data : 'a array,destination,tag) =    
        let arrType = typeof<'a> 
               
        match arrType with
            | Primitive -> mpiarrsend data (data.Length) (TTypeConvert arrType) destination tag
                                
            | Serializable 
             -> let arrLength = data.Length
                mpisend  arrLength  1 (TypeConvert arrLength)  destination tag
                let serializedArray = data |> Array.map (fun x -> serialize x)                              
                serializedArray |> Array.iter
                    (fun (str)  ->  mpiarrsend (str)  (int str.Length) (int MPI_Datatype.MPI_BYTE) destination tag)               
            | _ -> failwith "not supported from sendarray"                                
            
    let receiveArray<'a> (source, tag) =                             
        match typeof<'a>   with
        | Primitive ->  mpiarrayreceive<'a> source tag                        
                       
        | Serializable ->   let arrSize = mpireceive<int>  source tag 
                            let zArr = Array.zero_create  arrSize  :'a array
                            
                            zArr |> Array.map (fun i -> let buffer = mpiarrayreceive<byte> source tag
                                                        use memStream = new MemoryStream(buffer)
                                                        let desObj = deserialize buffer
                                                        memStream.Close(); 
  
                                                        desObj :?> 'a    )                                                                                    
                                                                 
        | _ -> failwith "Not supported from receivearray"
                
    let sendMatrix (data : matrix,destination,tag) =
        let  size = data.NumCols * data.NumRows        
        let pinnedMatrix = PinnedArray2.of_matrix data
        mpisend (data.NumRows) 1 (TypeConvert size) destination tag
        mpisend (data.NumCols) 1 (TypeConvert size) destination tag
        failonError  (MPI_Send( (pinnedMatrix.Ptr |> to_nativeint) ,size,(TTypeConvert (typeof<float>)) , destination , tag, MPI_COMM_WORLD)) "MPI_Send"                 
        pinnedMatrix.Free()        
   
    let receiveMatrix(source, tag) =        
        let row = mpireceive<int>    source tag
        let col = mpireceive<int>    source tag        
        let mutable status = new MPI_Status()
        let  data = Matrix.zero row col
        
        let pinnedM = PinnedArray2.of_matrix data                        
        failonError (MPI_Recv((pinnedM.Ptr |>  to_nativeint), (col*row), (TTypeConvert  (typeof<float>)), source, tag,MPI_COMM_WORLD, &&status)) "MPI_Recv"
        pinnedM.Free()
        data        
    
    let barrier() = failonError (MPI_Barrier(MPI_COMM_WORLD)) "MPI_Barrier"
    
    let packData (data:'a) = 
        let mutable data= data    
        match (box data).GetType() with
        | Primitive       ->  MPIChunk(&&data |> to_nativeint , 1 , (TypeConvert data))
        | Serializable _ -> let buffer= serialize data   
                            MPIChunkArray(buffer,buffer.Length, (int MPI_Datatype.MPI_BYTE))                            
        | _ -> failwith "Not supported from send function"   
                  
    let packArrayData (data:'a array) = 
       let arrType = typeof<'a>                
       match arrType with
            | Primitive -> MPIChunkArray(data,(data.Length),(TTypeConvert arrType) )
                                
            | Serializable 
             -> MPISerializedArray(data |> Array.map (fun x -> let str = serialize x in ((str), (int str.Length),(int MPI_Datatype.MPI_BYTE) ) ) )                
                
            | _ -> failwith "not supported from sendarray"              
              
    let broadcast(data,root)   =        
        match (packData data) with
         | MPIChunk (a,size,typ)     -> failonError (MPI_Bcast(a,size,typ,root, MPI_COMM_WORLD)) "MPI_Bcast"
         
         | MPIChunkArray(a,size,typ) -> let mutable size=size
                                        failonError (MPI_Bcast((&&size |> to_nativeint), 1,(TypeConvert size),root, MPI_COMM_WORLD)) "MPI_Bcast"
                                        let dat,frr = dataprepare a 
                                        
                                        try failonError (MPI_Bcast(dat,size,typ,root,MPI_COMM_WORLD)) "MPI_Bcast"
                                        finally freeA frr
         | _ -> ()
         
    let broadcastArray(data :'a array, root) =
        match (packArrayData data) with
           | MPIChunkArray(a,size,typ) ->   let mutable size=size
                                            failonError (MPI_Bcast((&&size |> to_nativeint), 1,(TypeConvert size),root, MPI_COMM_WORLD)) "MPI_Bcast"
                                            let dat,frr = dataprepare a 
                                        
                                            try failonError (MPI_Bcast(dat,size,typ,root,MPI_COMM_WORLD)) "MPI_Bcast"
                                            finally freeA frr
           
           | MPISerializedArray arr ->  let mutable size = arr.Length
                                        failonError (MPI_Bcast((&&size |> to_nativeint), 1,(TypeConvert size),root, MPI_COMM_WORLD)) "MPI_Bcast"
                                        
                                        arr |> Array.iter  (fun (a,size,typ) -> 
                                            let mutable size = size                                             
                                            failonError (MPI_Bcast((&&size |> to_nativeint), 1,(TypeConvert size),root, MPI_COMM_WORLD)) "MPI_Bcast" 
                                            
                                            let dat,frr = dataprepare a 
                                            try failonError (MPI_Bcast(dat, size,typ,root, MPI_COMM_WORLD)) "MPI_Bcast" 
                                            finally freeA frr                                                                                        
                                            )
                            
           | _ -> ()    
           
                 
    let private sCommon(data : 'a, root, count, f,fName) : 'a array = 
        let receiveArray :'a array= Array.zero_create count 
        match (packData data) with
         | MPIChunk (a,size,typ)     -> let dt,frr = dataprepare receiveArray
                                        try failonError (f(a,size,typ,dt, size,typ,root, MPI_COMM_WORLD)) fName
                                        finally freeA frr
                                        receiveArray
         | _ -> Array.zero_create 1
        
    let gather  (data:'a,root,count) :'a array= sCommon(data,root,count,MPI_Gather,"MPI_Gather")
    
    let scatter (data:'a,root,count) = sCommon(data,root,count,MPI_Scatter,"MPI_Scatter")
    
    let reduce(data,mpiop : MPI_OP, root) = 
        let receiveArray :'a array= Array.zero_create 1
        match (packData data) with
         | MPIChunk (a,size,typ)     -> let dt,frr = dataprepare receiveArray
                                        try failonError (MPI_Reduce(a,dt,size,typ,int mpiop,root, MPI_COMM_WORLD)) "MPI_Reduce"
                                        finally freeA frr
                                        receiveArray
         | _ -> Array.zero_create 1